use core::pin::Pin;
use std::future::Future;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};

/// Struct used for waiting on replies from the server.
///
/// When a BtlePlugMessage is sent to the server, it may take an indeterminate
/// amount of time to get a reply. This struct holds the reply, as well as a
/// [Waker] for the related future. Once the reply_msg is filled, the waker will
/// be called to finish the future polling.
#[derive(Debug, Clone)]
pub struct BtlePlugFutureState<T> {
    reply_msg: Option<T>,
    waker: Option<Waker>,
}

// For some reason, deriving default above doesn't work, but doing an explicit
// derive here does work.
impl<T> Default for BtlePlugFutureState<T> {
    fn default() -> Self {
        BtlePlugFutureState::<T> {
            reply_msg: None,
            waker: None,
        }
    }
}

impl<T> BtlePlugFutureState<T> {
    /// Sets the reply message in a message state struct, firing the waker.
    ///
    /// When a reply is received from (or in the in-process case, generated by)
    /// a server, this function takes the message, updates the state struct, and
    /// calls [Waker::wake] so that the corresponding future can finish.
    ///
    /// # Parameters
    ///
    /// - `msg`: Message to set as reply, which will be returned by the
    /// corresponding future.
    pub fn set_reply(&mut self, reply: T) {
        if self.reply_msg.is_some() {
            // TODO Can we stop multiple calls to set_reply_msg at compile time?
            panic!("set_reply_msg called multiple times on the same future.");
        }

        self.reply_msg = Some(reply);

        if self.waker.is_some() {
            self.waker.take().unwrap().wake();
        }
    }
}

/// Shared [BtlePlugFutureState] type.
///
/// [BtlePlugFutureState] is made to be shared across futures, and we'll
/// never know if those futures are single or multithreaded. Only needs to
/// unlock for calls to [BtlePlugFutureState::set_reply].
pub type BtlePlugFutureStateShared<T> = Arc<Mutex<BtlePlugFutureState<T>>>;

/// [Future] implementation for [BtlePlugMessageUnion] types send to the server.
///
/// A [Future] implementation that we can always expect to return a
/// [BtlePlugMessageUnion]. Used to deal with getting server replies after
/// sending [BtlePlugMessageUnion] types via the client API.
#[derive(Debug)]
pub struct BtlePlugFuture<T> {
    /// State that holds the waker for the future, and the [BtlePlugMessageUnion] reply (once set).
    ///
    /// ## Notes
    ///
    /// This needs to be an [Arc]<[Mutex]<T>> in order to make it mutable under
    /// pinning when dealing with being a future. There is a chance we could do
    /// this as a [Pin::get_unchecked_mut] borrow, which would be way faster, but
    /// that's dicey and hasn't been proven as needed for speed yet.
    waker_state: BtlePlugFutureStateShared<T>,
}

impl<T> Default for BtlePlugFuture<T> {
    fn default() -> Self {
        BtlePlugFuture::<T> {
            waker_state: BtlePlugFutureStateShared::<T>::default(),
        }
    }
}

impl<T> BtlePlugFuture<T> {
    /// Returns a clone of the state, used for moving the state across contexts
    /// (tasks/threads/etc...).
    pub fn get_state_clone(&self) -> BtlePlugFutureStateShared<T> {
        self.waker_state.clone()
    }

    // TODO Should we implement drop on this, so it'll yell if its dropping and
    // the waker didn't fire? otherwise it seems like we could have quiet
    // deadlocks.
}

impl<T> Future for BtlePlugFuture<T> {
    type Output = T;

    /// Returns when the [BtlePlugMessageUnion] reply has been set in the
    /// [BtlePlugFutureStateShared].
    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        let mut waker_state = self.waker_state.lock().unwrap();
        if waker_state.reply_msg.is_some() {
            let msg = waker_state.reply_msg.take().unwrap();
            Poll::Ready(msg)
        } else {
            waker_state.waker = Some(cx.waker().clone());
            Poll::Pending
        }
    }
}
